import numpy as np
import matplotlib.pyplot as plt

# 离散傅里叶变换
def dft(a, sign, opt=False):
    n = len(a)
    if a.dtype != np.complex:
        a = a.astype(np.complex)

    if opt:
        a[1: n+1: 2] = -a[1: n+1: 2]
    theta = np.linspace(0, 2*np.pi, n, endpoint=False)
    sin_tbl = sign * np.sin(theta)
    cos_tbl = np.cos(theta)

    b = np.zeros(n, dtype=np.complex)
    for i in range(n):
        for j in range(n):
            k = (i * j) % n
            c = np.complex(cos_tbl[k], sin_tbl[k])
            b[i] += a[j]*c

    if sign > 0:
        b *= 1.0 / n

    if opt:
        a[1: n+1: 2] = -a[1: n+1: 2]

    return b

# 测试傅里叶变换正确性
# a = np.array([0, 1, 2, 3, 4, 5, 6, 7], dtype=np.complex)
# b = np.fft.fft(a)
# c = dft(a, -1, True)
# np.set_printoptions(precision=4, suppress=True)
# print(b)
# print(c)

# 生成信号, 并进行傅里叶变换
x = np.linspace(0, 100, 1001)
#a = np.sin(x)
a = np.zeros(len(x), dtype=float)
for i in range(len(x)):
    if (i//50) % 2 == 0:
        a[i] = 1
b = dft(a, -1, True)
# 显示原始信号, 傅里叶振幅谱和相位谱
ax1 = plt.subplot(311)
ax1.set_title('Src signal')
plt.plot(x, a, 'r-')
ax2 = plt.subplot(312)
ax2.set_title('Amp spectrum')
plt.plot(x, np.abs(b), 'b-')
ax3 = plt.subplot(313)
ax3.set_title('Phase spectrum')
plt.plot(x, np.angle(b), 'g-')
plt.show()


